import os
import unittest
from autogen import AssistantAgent, UserProxyAgent

llm_config = {
  "config_list": [
    {
      "model": "qwen2:latest",
      "base_url": "http://127.0.0.1:11434/v1",
      "api_key": "qwen2_api_key",
      "cache_seed": None
    }
  ]
}

os.environ["http_proxy"] = "http://127.0.0.1:11434"  # 避免在开启 小猫时，无常访问 ollama 服务


class MyTestCase(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    cls.assistant = AssistantAgent(name="assistant", llm_config=llm_config)

    cls.user_proxy = UserProxyAgent(name="user_proxy", llm_config=llm_config,
                                    max_consecutive_auto_reply=0,
                                    human_input_mode="NEVER",
                                    code_execution_config={
                                      "work_dir": "coding",
                                      "use_docker": False
                                    })
    # 在这里进行一次性的初始化操作
    print("setUpClass executed")

  def notest_something(self):
    self.user_proxy.initiate_chat(self.assistant, message="中国的首都是哪里，到南极大约多少公里")

  def test_down_docs(self):
    message = """
    Barsoum RS. Chronic kidney disease in the developing world. N Engl J Med 2006; 354: 997-9. 
    
    请根据以上引用信息，查询它的DOI 信息，请使用 JSON格式返回最终结果 。
    """
    message2="""
任务: 根据论文引用查询 DOI 信息

引用信息:
- 作者: Barsoum RS.
- 标题: Chronic kidney disease in the developing world
- 期刊名称: N Engl J Med
- 期号: 354
- 页码: 997-999
- 出版年份: 2006
    """
    self.user_proxy.initiate_chat(self.assistant, message=message)


if __name__ == '__main__':
  unittest.main()
